
/* 接收端：
1. 订阅包含data的mqtt主题
2. 从另一个来源获取元数据
3. 创建一个新的 sensor_msgs::msg::PointCloud2 消息。
4. 将获取的元数据填充到新消息的对应字段。
5. 将从 MQTT 接收到的 data 字节数组复制到新消息的 data 字段。
6. 发布重构的 PointCloud2 消息到 ROS 话题。 */

#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/msg/point_cloud2.hpp"
#include <mosquitto.h>
#include <jsoncpp/json/json.h>
#include <string>
#include <vector>
#include <memory>
#include <cstring> 

class PointCloudMqttSubscriber : public rclcpp::Node
{
public:
    PointCloudMqttSubscriber() : Node("pointcloud_mqtt_subscriber")
    {
        // 声明参数
        mqtt_host_ = this->declare_parameter<std::string>("mqtt_net.host", "localhost");
        mqtt_port_ = this->declare_parameter<int>("mqtt_net.port", 1883);
        output_ros_topic_ = this->declare_parameter<std::string>("ros.output_topic", "/mqtt_reconstructed_cloud");
        mqtt_data_topic_ = this->declare_parameter<std::string>("mqtt.data_topic", "/fangshe/robot/lidarb/cloud/data");
        mqtt_meta_topic_ = this->declare_parameter<std::string>("mqtt.meta_topic", "/fangshe/robot/lidarb/cloud/meta");

        // 获取参数值
        mqtt_host_ = this->get_parameter("mqtt_net.host").as_string();
        mqtt_port_ = this->get_parameter("mqtt_net.port").as_int();
        output_ros_topic_ = this->get_parameter("ros.output_topic").as_string();
        mqtt_data_topic_ = this->get_parameter("mqtt.data_topic").as_string();
        mqtt_meta_topic_ = this->get_parameter("mqtt.meta_topic").as_string();

        RCLCPP_INFO(this->get_logger(), "Output ROS Topic: %s", output_ros_topic_.c_str());
        RCLCPP_INFO(this->get_logger(), "MQTT Data Topic: %s", mqtt_data_topic_.c_str());
        RCLCPP_INFO(this->get_logger(), "MQTT Meta Topic: %s", mqtt_meta_topic_.c_str());

        // 初始化 MQTT
        mosquitto_lib_init();
        mqtt_client_ = mosquitto_new(nullptr, true, this);

        if (!mqtt_client_)
        {
            RCLCPP_ERROR(this->get_logger(), "Failed to create MQTT client");
            return;
        }

        mosquitto_message_callback_set(mqtt_client_, on_message_static);
        mosquitto_connect_callback_set(mqtt_client_, on_connect_static);
        mosquitto_disconnect_callback_set(mqtt_client_, on_disconnect_static);

        connect_thread_ = std::thread(&PointCloudMqttSubscriber::connect_loop, this);

        int rc = mosquitto_loop_start(mqtt_client_);
        if (rc != 0)
        {
            RCLCPP_ERROR(this->get_logger(), "Could not start MQTT loop: %s", mosquitto_strerror(rc));
            mosquitto_disconnect(mqtt_client_);
            mosquitto_destroy(mqtt_client_);
            return;
        }

        // 创建 ROS 发布者
        pointcloud_pub_ = this->create_publisher<sensor_msgs::msg::PointCloud2>(output_ros_topic_, 10);

        RCLCPP_INFO(this->get_logger(), "PointCloud MQTT Subscriber initialized.");
    }

    ~PointCloudMqttSubscriber()
    {
        stop_connect_thread_ = true;
        if (connect_thread_.joinable())
        {
            connect_thread_.join();
        }

        if (mqtt_client_)
        {
            mosquitto_loop_stop(mqtt_client_, true);
            mosquitto_disconnect(mqtt_client_);
            mosquitto_destroy(mqtt_client_);
        }
        mosquitto_lib_cleanup();
    }

private:
    void on_message(const struct mosquitto_message *msg)
    {
        std::string topic(msg->topic);

        if (topic == mqtt_meta_topic_)
        {
            RCLCPP_DEBUG(this->get_logger(), "Received metadata on [%s], size: %zu", topic.c_str(), msg->payloadlen);

            // 解析 JSON 元数据
            Json::Value root;
            Json::CharReaderBuilder reader_builder;
            std::string errors;
            std::unique_ptr<Json::CharReader> reader(reader_builder.newCharReader());
            std::string payload_str(static_cast<const char*>(msg->payload), msg->payloadlen);

            if (reader->parse(payload_str.c_str(), payload_str.c_str() + payload_str.length(), &root, &errors))
            {
                // 将解析后的元数据存储起来
                cached_metadata_.header.stamp.sec = root["header"]["stamp"]["sec"].asInt();
                cached_metadata_.header.stamp.nanosec = root["header"]["stamp"]["nanosec"].asUInt();
                cached_metadata_.header.frame_id = root["header"]["frame_id"].asString();
                cached_metadata_.height = root["height"].asUInt();
                cached_metadata_.width = root["width"].asUInt();
                cached_metadata_.is_bigendian = root["is_bigendian"].asBool();
                cached_metadata_.point_step = root["point_step"].asUInt();
                cached_metadata_.row_step = root["row_step"].asUInt();
                cached_metadata_.is_dense = root["is_dense"].asBool();

                cached_metadata_.fields.clear();
                const Json::Value& fields_json = root["fields"];
                for (const auto& field_json : fields_json)
                {
                    sensor_msgs::msg::PointField field;
                    field.name = field_json["name"].asString();
                    field.offset = field_json["offset"].asUInt();
                    field.datatype = static_cast<uint8_t>(field_json["datatype"].asInt());
                    field.count = field_json["count"].asUInt();
                    cached_metadata_.fields.push_back(field);
                }

                RCLCPP_DEBUG(this->get_logger(), "Cached metadata for frame_id: %s", cached_metadata_.header.frame_id.c_str());
            }
            else
            {
                RCLCPP_ERROR(this->get_logger(), "Failed to parse metadata JSON: %s", errors.c_str());
            }
        }
        else if (topic == mqtt_data_topic_)
        {
            RCLCPP_INFO(this->get_logger(), "Received data on [%s], size: %zu", topic.c_str(), msg->payloadlen);

            // 检查是否有缓存的元数据
            if (cached_metadata_.header.frame_id.empty())
            {
                RCLCPP_WARN(this->get_logger(), "Received data but no cached metadata available. Dropping data packet.");
                return; // 没有元数据无法重构消息
            }

            // 创建新的 PointCloud2 消息
            sensor_msgs::msg::PointCloud2::SharedPtr reconstructed_msg = std::make_shared<sensor_msgs::msg::PointCloud2>();

            // 填充元数据
            reconstructed_msg->header = cached_metadata_.header;
            reconstructed_msg->height = cached_metadata_.height;
            reconstructed_msg->width = cached_metadata_.width;
            reconstructed_msg->fields = cached_metadata_.fields;
            reconstructed_msg->is_bigendian = cached_metadata_.is_bigendian;
            reconstructed_msg->point_step = cached_metadata_.point_step;
            reconstructed_msg->row_step = cached_metadata_.row_step;
            reconstructed_msg->is_dense = cached_metadata_.is_dense;

            // 填充二进制数据
            reconstructed_msg->data.resize(msg->payloadlen);
            std::memcpy(reconstructed_msg->data.data(), msg->payload, msg->payloadlen);

            // 发布重构的消息
            pointcloud_pub_->publish(*reconstructed_msg);
            RCLCPP_INFO(this->get_logger(), "Published reconstructed PointCloud2 to ROS topic: %s, size: %zu bytes", 
                        output_ros_topic_.c_str(), reconstructed_msg->data.size());

            // 清除缓存的元数据 (如果元数据不频繁变化，可以不清除，或者用时间戳/序列号匹配)
            // cached_metadata_ = MetadataCache(); // 重置结构
        }
    }

    // --- MQTT Callbacks (Static) ---
    static void on_message_static(struct mosquitto *mosq, void *obj, const struct mosquitto_message *msg)
    {
        if (!obj || !msg) return;
        static_cast<PointCloudMqttSubscriber*>(obj)->on_message(msg);
    }

    static void on_connect_static(struct mosquitto *mosq, void *obj, int rc)
    {
        if (!obj) return;
        static_cast<PointCloudMqttSubscriber*>(obj)->on_connect(rc);
    }

    void on_connect(int rc)
    {
        if (rc == 0)
        {
            is_connected_ = true;
            RCLCPP_INFO(this->get_logger(), "Successfully connected to MQTT broker");

            // 连接成功后订阅主题
            resubscribe_topics();
        }
        else
        {
            is_connected_ = false;
            RCLCPP_ERROR(this->get_logger(), "Failed to connect to MQTT broker: %s", mosquitto_strerror(rc));
        }
    }

    static void on_disconnect_static(struct mosquitto *mosq, void *obj, int rc)
    {
        if (!obj) return;
        static_cast<PointCloudMqttSubscriber*>(obj)->on_disconnect(rc);
    }

    void on_disconnect(int rc)
    {
        if (rc != 0)
        {
            is_connected_ = false;
            RCLCPP_WARN(this->get_logger(), "Disconnected from MQTT broker. Will attempt to reconnect...");
        }
    }

    void connect_loop()
    {
        int retry_count = 0;
        const int max_retry_delay = 5;

        while (rclcpp::ok() && !stop_connect_thread_)
        {
            if (!is_connected_)
            {
                RCLCPP_INFO(this->get_logger(), "Trying to connect to MQTT broker: %s:%d",
                            mqtt_host_.c_str(), mqtt_port_);

                int rc = mosquitto_connect(mqtt_client_, mqtt_host_.c_str(), mqtt_port_, 60);
                if (rc != 0)
                {
                    int delay = std::min(1 << retry_count, max_retry_delay);
                    RCLCPP_ERROR(this->get_logger(), "Could not connect to MQTT broker: %s. Retrying in %d seconds...",
                                 mosquitto_strerror(rc), delay);
                    std::this_thread::sleep_for(std::chrono::seconds(delay));
                    retry_count++;
                }
                else
                {
                    std::this_thread::sleep_for(std::chrono::seconds(1));
                }
            }
            else
            {
                std::this_thread::sleep_for(std::chrono::seconds(1));
            }
        }
    }

    void resubscribe_topics()
    {
        std::vector<std::string> topics = {mqtt_data_topic_, mqtt_meta_topic_};

        for (const auto &topic : topics)
        {
            int rc = mosquitto_subscribe(mqtt_client_, nullptr, topic.c_str(), 0);
            if (rc != 0)
            {
                RCLCPP_ERROR(this->get_logger(), "Could not subscribe to MQTT topic [%s]: %s", topic.c_str(), mosquitto_strerror(rc));
            }
            else
            {
                RCLCPP_INFO(this->get_logger(), "Subscribed to MQTT topic: %s", topic.c_str());
            }
        }
    }

    // 用于缓存元数据的结构
    struct MetadataCache
    {
        std_msgs::msg::Header header;
        uint32_t height;
        uint32_t width;
        std::vector<sensor_msgs::msg::PointField> fields;
        bool is_bigendian;
        uint32_t point_step;
        uint32_t row_step;
        bool is_dense;

        // 默认构造函数
        MetadataCache() : height(0), width(0), is_bigendian(false), point_step(0), row_step(0), is_dense(false) {}
    };

    struct mosquitto *mqtt_client_;
    std::string output_ros_topic_;
    std::string mqtt_data_topic_;
    std::string mqtt_meta_topic_;

    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr pointcloud_pub_;

    bool is_connected_ = false; // Initialize
    std::thread connect_thread_;
    bool stop_connect_thread_ = false;
    std::string mqtt_host_;
    int mqtt_port_;

    // 缓存最后接收到的元数据
    MetadataCache cached_metadata_;
};

int main(int argc, char *argv[])
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<PointCloudMqttSubscriber>());
    rclcpp::shutdown();
    return 0;
}